import pandas as pd
from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.agents.format_scratchpad.tools import format_to_tool_messages
from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_experimental.tools import PythonAstREPLTool

load_dotenv(override=True)



#定义模型
model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
df = pd.read_csv("./datas/WA_Fn-UseC_-Telco-Customer-Churn.csv")
pd.set_option("max_colwidth", 200)
pythonREPL = PythonAstREPLTool(locals={"df":df})

#基础功能测试
def dataTest():
    print(df.head(5))
    print(pythonREPL.invoke("df['SeniorCitizen'].mean()"),",",df['SeniorCitizen'].mean())
    print(pythonREPL.invoke("df['MonthlyCharges'].mean()"),",",df['MonthlyCharges'].mean())

dataTest()

# 定义 天气查询 工具函数
tools = [pythonREPL]
llm_with_tools = model.bind_tools(tools)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是数据处理专家，善于帮助用户处理数据方面的工作"),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)

# agent = (
#     RunnablePassthrough.assign(
#         agent_scratchpad=lambda x: format_to_tool_messages(x["intermediate_steps"])
#     )
#     | prompt
#     | llm_with_tools
#     | ToolsAgentOutputParser()
# )

agent = create_tool_calling_agent(llm=model,tools=tools,prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

while True:
    query = input("请问我可以帮您做什么吗？")

    if query == "exit" or query == "quit" or query == "bye" or query == "goodbye" or query == "stop":
        break
    #并行
    response = agent_executor.invoke({"input": query})
    print(response['output'])
